open Ast
open Http_client.Convenience
module StringMap = Map.Make (String);;

type symbol_table = {
  parent : symbol_table option; 
  vars : lit StringMap.t;
  funcs : (((datatype*string) list)*stmt) StringMap.t;
}

type translation_environment = {
    return_type : datatype;     (* Function's return value *)
    in_loop : bool;             (* whether break and continue are valid *)
    scope : symbol_table;       (* symbol table for vars *)
  }

exception TypeError
exception ReturnCall of lit*symbol_table
exception BreakStatement of symbol_table
exception Fatal of string

(* Main entry point: run a program *)
let run (stmts) = 

  (* initial symbol table *)
  let empty_symbol = {
    parent = None;
    vars = StringMap.empty;
    funcs = StringMap.empty;}
  in

  (* Perform variable lookup *)
  let rec findvar env = function
      IdVar(name) ->
        try (StringMap.find name env.vars)
          with Not_found -> match env.parent with
                                None -> raise Not_found
                              | Some p -> findvar p (IdVar name)
  in

  (* Perform a function lookup *)
  let rec findfunc env = function
      name ->
        try (StringMap.find name env.funcs)
          with Not_found -> match env.parent with
                                None -> raise Not_found
                              | Some p -> findfunc p name
  in

  (* find the scope of a var and set it *)
  let rec setvar env = function
      IdVar(name), v ->
        try (let _ = StringMap.find name env.vars in
             v, {parent = env.parent;
                 vars = StringMap.add name v env.vars;
                 funcs = env.funcs;})
          with Not_found -> match env.parent with
                                None -> raise Not_found
                              | Some p ->
                                  let v, p = setvar p ((IdVar name), v) in
                                  v, {parent = Some p;
                                      vars = env.vars;
                                      funcs = env.funcs;}
  in

  (* find the scope of an array and set an element *)
  let rec sete i v l = (i, v) :: (List.remove_assoc i l) in
  let rec setmap env = function
      IdVar(name), i, v ->
        (try (match StringMap.find name env.vars with
                  Literal_map(old) ->
                    v, {parent = env.parent;
                        vars = StringMap.add name
                                 (Literal_map (sete i v old)) env.vars;
                        funcs = env.funcs;}
                | _ -> raise TypeError)
           with Not_found -> match env.parent with
                                 None -> raise Not_found
                               | Some p ->
                                   let v, p = setmap p
                                                ((IdVar name), i, v) in
                                   v, {parent = Some p;
                                       vars = env.vars;
                                       funcs = env.funcs;})
  in

  (* find the scope of a map and set an element *)
  let rec setn i v = function
      hd :: rest ->
        if (i < 0) then (hd :: rest) else
          if (i = 0) then (v :: rest) else
            (hd :: (setn (i - 1) v rest))
    | [] -> []
  in 
  let rec setarr env = function
      IdVar(name), Literal_int(i), v ->
        (try (match StringMap.find name env.vars with
                  Literal_array(old) ->
                    v, {parent = env.parent;
                        vars = StringMap.add name
                                 (Literal_array (setn i v old)) env.vars;
                        funcs = env.funcs;}
                | _ -> raise TypeError)
           with Not_found -> match env.parent with
                                 None -> raise Not_found
                               | Some p ->
                                   let v, p = setarr p
                                                ((IdVar name),
                                                 (Literal_int i),
                                                 v) in
                                   v, {parent = Some p;
                                       vars = env.vars;
                                       funcs = env.funcs;})
    | _ -> raise TypeError
  in

  (* Check if a type and literal match *)
  let rec check_type dtype lit =
    (match dtype,lit with
         Type_int,Literal_int(a) -> true
       | Type_bool,Literal_bool(a) -> true
       | Type_float,Literal_float(a) -> true
       | Type_str,Literal_str(a) -> true
       | Type_Array(dtype),
         Literal_array(hd::rest) -> check_type dtype hd
       | Type_Array(dtype), Literal_array([]) -> true
       | Type_Map(ktype, vtype),
         Literal_map((klit, vlit)::rest) ->
           ((check_type ktype klit) && (check_type vtype vlit))
       | Type_html,Literal_html(_) -> true
       | _ -> false)
  in

  (* Evaluate an expression and return (value, updated environment) *)
  let rec eval env = function
      Literal(i) ->  i, env
    | Arrayref(name, e) ->
        let v, env = eval env e in
        let arr = findvar env name in
        (match v, arr with
             Literal_int(i), Literal_array(l) -> List.nth l i
           | _ -> raise TypeError), env
    | Mapref(name, e) ->
        let v, env = eval env e in
        let map = findvar env name in
        (match map with
             Literal_map(mapcontent) ->
              (List.assoc v mapcontent)
           | _ -> raise TypeError), env
    | Arrayassign(name, e1, e2) -> 
        let v1, env = eval env e1 in
        let v2, env = eval env e2 in
        setarr env (name, v1, v2)
    | Mapassign(name, e1, e2) ->
        let v1, env = eval env e1 in
        let v2, env = eval env e2 in
        setmap env (name, v1, v2)
    | Var(var) -> (findvar env var), env
    | Binop(e1, op, e2) ->
        let l1, env = eval env e1 in
        let l2, env = eval env e2 in
        (match (l1, l2) with
          (Literal_int(v1), Literal_int(v2)) ->
            (match op with
              Add -> Literal_int(v1 + v2)
            | Sub -> Literal_int(v1 - v2)
            | Mult -> Literal_int(v1 * v2)
            | Div -> Literal_int(v1 / v2)
            | Equal -> Literal_bool(v1 = v2)
            | Neq -> Literal_bool(v1 != v2)
            | Less -> Literal_bool(v1 < v2)
            | Leq -> Literal_bool(v1 <= v2)
            | Greater -> Literal_bool(v1 > v2)
            | Geq -> Literal_bool(v1 >= v2)
            | _ -> raise TypeError)
        | (Literal_str(v1), Literal_str(v2)) ->
            (match op with
                 Add -> Literal_str(v1 ^ v2)
               | Equal -> Literal_bool(v1 = v2)
               | Neq -> Literal_bool(v1 != v2)
               | _ -> raise TypeError)
        | _ -> raise TypeError), env
    | Uniop(op, e) -> 
        let l, env = eval env e in
          (match l with
               Literal_bool(v) ->
                 (match op with
                    Not -> Literal_bool(not v))
             | _ -> raise TypeError), env
    | Assign(IdVar(varname), e) ->
        let v, env = eval env e in
        setvar env ((IdVar varname), v)

(* Start of Libraries *)
  (* String Library *)
    | Call("substring", [e1; e2; e3]) -> 
        let v1, env = eval env e1 in
        let v2, env = eval env e2 in
        let v3, env = eval env e3 in
        (match [v1;v2;v3] with [Literal_str(v1); Literal_int(v2); Literal_int(v3)]  -> (
               try 
                 Literal_str( String.sub v1 v2 ((v3 - v2)) )
               with (Invalid_argument(z)) -> raise (Invalid_argument(z)) )
        | _ -> raise TypeError), env

    | Call("stringLength", [e1]) -> 
        let v1, env = eval env e1 in
        (match v1 with Literal_str(v1)  -> 
          Literal_int( String.length v1 );
        | _ -> raise TypeError), env

    | Call("stringFind", [e1;e2]) ->	 
        let v1, env = eval env e1 in
	let v2, env = eval env e2 in
        (match [v1;v2] with [Literal_str(v1); Literal_str(v2)]  ->
          Literal_int( try Str.search_forward (Str.regexp v2) v1 0 with Not_found->0 );
        | _ -> raise TypeError), env		
    | Call("stringReplace", [e1;e2;e3]) -> 
        let v1, env = eval env e1 in
	let v2, env = eval env e2 in
	let v3, env = eval env e3 in
        (match [v1;v2;v3] with [Literal_str(v1); Literal_str(v2);Literal_str(v3)]  ->
          Literal_str(Str.global_replace (Str.regexp v1) v2 v3);
        | _ -> raise TypeError), env
		
    | Call("stringToUpper", [e1]) -> 
        let v1, env = eval env e1 in
        (match v1 with Literal_str(v1)  ->
          Literal_str( String.uppercase v1);
        | _ -> raise TypeError), env
		
    | Call("stringToLower", [e1]) -> 
        let v1, env = eval env e1 in
        (match v1 with Literal_str(v1)  ->
          Literal_str( String.lowercase v1);
        | _ -> raise TypeError), env

    | Call("stringFindAll", [e1; e2]) -> 
        let v1, env = eval env e1 in
        let v2, env = eval env e2 in
        (match [v1;v2] with 
          [Literal_str(v1); Literal_str(v2)]  -> 
               let rec stringFindAllHelper inputString searchString index indexList = 
                   try 
                      let currIndex = Str.search_forward (Str.regexp searchString) inputString index in 
                      let nextIndex = (currIndex + (String.length searchString)) in
                      let currIndex2 = Literal_int(currIndex) in
                      (stringFindAllHelper inputString searchString nextIndex  (currIndex2::indexList)) 
                   with Not_found -> indexList; 
               in
               Literal_array (List.rev(stringFindAllHelper v1 v2 0 []));
        | _ -> raise TypeError), env

    | Call("splitString", [e1;e2]) -> 
        let v1, env = eval env e1 in
	let v2, env = eval env e2 in
        (match [v1;v2] with [Literal_str(v1); Literal_str(v2)]  ->
          Literal_array( List.map (fun f-> Literal_str f) (Str.split (Str.regexp v1) v2) );
        | _ -> raise TypeError), env

  (* Array Library *)
      (* Works for array and maps *)
    | Call("length", [e1]) -> 
        let v1, env = eval env e1 in
        (match v1 with 
          Literal_array(v1) -> Literal_int(List.length v1)
        | Literal_map(v1)   -> Literal_int(List.length v1)
        | _ -> raise TypeError), env	

      (* Maps are represented as (lit*lit) lists *)
    | Call("getKeys", [e1]) -> 
        let v1, env = eval env e1 in
        (match v1 with Literal_map(v1)  ->
          Literal_array( List.rev( List.fold_left (fun l (a,b)->a::l) [] v1) );
        | _ -> raise TypeError), env	
		
    | Call("getValues", [e1]) -> 
        let v1, env = eval env e1 in
        (match v1 with Literal_map(v1)  ->
          Literal_array( List.rev( List.fold_left (fun l (a,b)->b::l) [] v1) );
        | _ -> raise TypeError), env	

      (* Works for array and maps according to key, e2 true for ascending *)
    | Call("sort", [e1;e2]) -> 
        let v1, env = eval env e1 in
	let v2, env = eval env e2 in
        (match [v1;v2] with 
          [Literal_array(v1); Literal_bool(v2)] -> Literal_array( if v2 then (List.sort compare v1)
                                                                        else List.rev (List.sort compare v1) )
        | [Literal_map(v1);   Literal_bool(v2)] -> Literal_map ( let x = List.sort (fun (a,b) (c,d) -> if a<c then -1 else
                                                                                   (if a=c then 0 else 1)) v1
                                                                 in 
                                                                 if v2 then x else List.rev x )
        | _ -> raise TypeError), env
		
    | Call("randomize", [e1]) -> 
        let v1, env = eval env e1 in
        (match v1 with
          Literal_array(v1) -> let _ = Random.self_init() in (* DO NOT REMOVE, REQUIRED FOR RANDOM #s *)
                               Literal_array( (List.sort (fun x y -> Random.int 2) v1) );
        | _ -> raise TypeError), env

    | Call("numargs", []) ->
        Literal_int(Array.length Sys.argv), env
    | Call("getarg", [e1]) ->
        let v1, env = eval env e1 in
        (match v1 with
            Literal_int(v1) -> Literal_str(Sys.argv.(v1))
          | _ -> raise TypeError), env

     (*  HTML Library   *)
    | Call("url", [e1]) ->
        let v1, env = eval env e1 in
        ( match v1 with Literal_str(v1) -> Literal_array ( List.map (fun f->Literal_html(f)) (Nethtml.parse_document (Lexing.from_string (http_get v1))))
        | _ -> raise TypeError), env

    | Call("filterArrayByAttribute", [e1;e2;e3]) ->
        let v1, env = eval env e1 in
        let v2, env = eval env e2 in
        let v3, env = eval env e3 in
        ( match [v1;v2;v3] with [Literal_str(v1);Literal_str(v2);Literal_array(v3)] ->
          let filterByAttribute attr value attributeMap =
            if ( List.mem_assoc attr attributeMap)
             then if ((List.assoc attr attributeMap) = value)
               then [Literal_map(attributeMap)]
               else []
             else [] in
          let strippedMaps = (List.map (function Literal_map(x) -> x
                                               | _ -> raise TypeError) v3) in
          let filteredMaps = (List.fold_left
                               (fun retList am -> List.append retList (filterByAttribute (Literal_str v1) (Literal_str v2) am)) [] strippedMaps) in
          Literal_array(filteredMaps)
        | _ -> raise TypeError), env

   | Call("getArrayByTag", [e1;e2]) ->
        let v1, env = eval env e1 in
        let v2, env = eval env e2 in
        (match [v1;v2] with [Literal_str(l1);Literal_array(l2)] ->
               let rec find tag = (function Nethtml.Element(content,a,c) -> if (tag=content) then [(List.map (function (x,y) -> Literal_str(x), Literal_str(y))a)] else List.concat(List.map (find tag) c)
                                          | Nethtml.Data(e) -> []) in
               let getListByTag tag docList = (List.concat (List.map (find tag) docList)) in
               Literal_array(List.map (fun a -> Literal_map(a)) (getListByTag l1 (List.map (function Literal_html(x) -> x | _ -> raise TypeError) l2)))
               | _ -> raise TypeError), env 


(* End of Libraries *)

    | Call("print", [e]) ->
        let v, env = eval env e in
        (match v with
          Literal_int(intval) -> 
            print_endline (string_of_int intval);
            Literal_int(0)
        | Literal_bool(boolval) -> 
            print_endline (string_of_bool boolval);
            Literal_int(0)
        | Literal_float(floatval) -> 
            print_endline (string_of_float floatval);
            Literal_int(0)
        | Literal_str(strval) -> 
            let rec f s =
              try let x = String.index s '\\' in
                  print_string((String.sub s 0 x)^(String.sub s (x+1) 1));
                  f((String.sub s (x+2) ((String.length s)-(x+2))));
              with Not_found -> print_endline(s); in
            f(strval);
            Literal_int(0)
        | _ -> raise TypeError), env

  (* Print without ending line *)
| Call("inprint", [e]) ->
        let v, env = eval env e in
        (match v with
          Literal_int(intval) -> 
            print_string (string_of_int intval);
            Literal_int(0)
        | Literal_bool(boolval) -> 
            print_string (string_of_bool boolval);
            Literal_int(0)
        | Literal_float(floatval) -> 
            print_string (string_of_float floatval);
            Literal_int(0)
        | Literal_str(strval) -> 
            let rec f s =
              try let x = String.index s '\\' in
                  print_string((String.sub s 0 x)^(String.sub s (x+1) 1));
                  f((String.sub s (x+2) ((String.length s)-(x+2))));
              with Not_found -> print_string(s); in
            f(strval);
            Literal_int(0)
        | _ -> raise TypeError), env

    | Call(func, args) ->
        let exp_args, body = findfunc env func in
        let actuals, env = List.fold_left
            (fun (actuals, env) arg ->
              let v, env = eval env arg in
                (v :: actuals, env)) ([], env) args in
        let topass = List.combine exp_args actuals in
        let newenv = {
          parent = Some env;
          vars = List.fold_left
                   (fun map ((_, argname), argval) ->
                     StringMap.add argname argval map) StringMap.empty topass;
          funcs = StringMap.empty; } in
        let v, newenv = try Literal_int(0), (exec newenv body)
                          with ReturnCall(v,env) -> v, env in
        (match newenv.parent with
             Some p -> v, p
          | _ -> raise (Fatal "Function call destroyed environment"))
    | Noexpr -> Literal_int(0), env

  (* Execute a statement and return an updated environment *)
  and exec env = function
      Block(stmts) -> List.fold_left exec env stmts
    | Expr(e) -> let _, env = eval env e in env
    | Return(e) -> let v, env = eval env e in raise (ReturnCall(v,env))
    | Use(name) -> env
    | Break -> raise (BreakStatement env)
    | If(e, s1, s2) ->
        let v, env = eval env e in
        exec env (if v = Literal_bool(true) then s1 else s2)
    | Foreach(id1,id2,e,s) ->
        let v, env = eval env e in
        (match v with
             Literal_array(l) ->
               let size = List.length l in
               let rec loop env i = function
                   hd :: rest ->
                     let _, env = setvar env (IdVar(id1), Literal_int(i)) in
                     let _, env = setvar env (IdVar(id2), hd) in
                     if i >= size then
                       env
                     else
                       loop (exec env s) (i + 1) rest
                 | [] -> env in
               (try (loop env 0 l) with BreakStatement(env) -> env)
           | Literal_map(l) ->
               let rec loop env = function
                   (kval, vval) :: rest ->
                      let _, env = setvar env (IdVar(id1), kval) in
                      let _, env = setvar env (IdVar(id2), vval) in
                      loop (exec env s) rest
                 | [] -> env in
               (try (loop env l) with BreakStatement(env) -> env)
           | _ -> raise TypeError)
    | While(e, s) ->
        let rec loop env =
          let v, env = eval env e in
          if v = Literal_bool(true) then loop (exec env s) else env
        in (try (loop env) with BreakStatement(env) -> env)
    | Declare(dtype, name) -> 
        let rec makedefault = function
               Type_int -> Literal_int(0)
             | Type_bool -> Literal_bool(true)
             | Type_float -> Literal_float(0.0)
             | Type_str -> Literal_str("")
             | Type_Array(dtype) -> Literal_array([])
             | Type_Map(ktype, ltype) ->
                (Literal_map [])
             | Type_html -> Literal_html(Nethtml.Data("")) in
        let nv = (makedefault dtype) in
          {parent = env.parent;
           vars = StringMap.add name nv env.vars;
           funcs = env.funcs;}
    | DeclareAssign(dtype, name, e) ->
        let v, env = eval env e in
        if (check_type dtype v) then
          {parent = env.parent;
           vars = StringMap.add name v env.vars;
           funcs = env.funcs;}
        else
          raise TypeError
    | FuncDecl(name, args, rettype, body) ->
        {parent = env.parent;
         vars = env.vars;
         funcs = StringMap.add name (args, body) env.funcs;}
  in
  try (List.fold_left exec empty_symbol stmts)
    with
        TypeError ->
          print_endline "Error: type mismatch";
          empty_symbol
      | ReturnCall(_) ->
          print_endline "Error: return from outside function body";
          empty_symbol
      | Invalid_argument(_) -> 
          print_endline "Invalid arguments";
          empty_symbol
      | Not_found ->
          print_endline "Variable does not exist";
          empty_symbol
